﻿using ELFSharp.ELF;
using MySoftLab.DapFlashCore.CoreSight;
using MySoftLab.DapFlashCore.Probes.DapAccess;
using MySoftLab.DapFlashCore.Targets;
using MySoftLab.DapFlashCore.Utils;
using System;
using System.Collections;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Text;
using System.Threading;
using System.Threading.Tasks;

namespace MySoftLab.DapFlashCore.Flash
{
    public class FlashMemory
    {
        private readonly CortexM _cpu;
        private FlashOperations _activeOperation = FlashOperations.NONE;
        private readonly FlashAlgo _flashAlgo;
        private bool _didPrepareTarget = false;
        private uint _endFlashAlgo;

        public const int TIMEOUT_ERROR = -1;

        public bool IsEraseAllSupported
        {
            get
            {
                return IsApiValid(_flashAlgo.PcEraseAll);
            }
        }

        private readonly TargetBase _target;

        // 进度回调委托定义
        public delegate void ProgressCallbackDelegate(int current, int total);
        
        // 进度回调属性
        public ProgressCallbackDelegate? ProgressCallback { get; set; }

        // 添加日志委托定义
        public delegate void LogCallbackDelegate(string message);
        public LogCallbackDelegate? AppendLog { get; set; }

        public FlashMemory(CortexM cpu, TargetBase target)
        {
            _cpu = cpu;
            _target = target;
            _flashAlgo = target.FlashAlgo;

            // TODO 初始化各个值
            _endFlashAlgo = (uint)(_flashAlgo.LoadAddress + _flashAlgo.Instructions.Length * 4);
        }

        public void Init(FlashOperations operation, uint address = 0, uint clock = 0, bool reset = true)
        {
            // 首先暂停target
            _cpu.Halt();

            if (address == 0)
            {
                address = 0x08000000;
            }

            if (_activeOperation != FlashOperations.NONE)
            {
                if (_activeOperation != operation)
                {
                    UnInit();
                }
                else
                {
                    return;
                }
            }

            // 运行flash算法之前初始化目标
            if (!_didPrepareTarget)
            {
                if (reset)
                {
                    _cpu.ResetAndHalt();
                }

                PrepareTarget();

                // 加载flash算法至目标RAM
                _cpu.WriteMemoryBlock32(_flashAlgo.LoadAddress, _flashAlgo.Instructions);

                // var data = _cpu.ReadMemoryBlock32(_flashAlgo.LoadAddress, _flashAlgo.Instructions.Length);

                _didPrepareTarget = true;
            }

            // update core register to execute the init subroutine
            if (IsApiValid(_flashAlgo.PcInit))
            {
                Console.WriteLine($"algo call init(addr={address}, clock={clock}, op={operation})");

                int result = CallFunctonAndWait(_flashAlgo.PcInit, r0: address, r1: clock, r2: (uint)operation, init: true, timeout: 5);
                // check the return code
                Console.WriteLine($"init result = {result}");
                if (result == TIMEOUT_ERROR)
                {
                    throw new Exception("flash init time out");
                }
                else if (result != 0)
                {
                    throw new Exception($"flash init failure, code = {result}");
                }
            }

            _activeOperation = operation;
        }

        private bool IsApiValid(uint address)
        {
            return address < _endFlashAlgo && address >= _flashAlgo.LoadAddress;
        }

        public void UnInit()
        {
            if (_activeOperation == FlashOperations.NONE)
            {
                return;
            }

            if (IsApiValid(_flashAlgo.PcUnInit))
            {
                Console.WriteLine($"call uninit({_activeOperation})");

                // update core register to execute the uninit subroutine
                var result = CallFunctonAndWait(_flashAlgo.PcUnInit, r0: (uint)_activeOperation, timeout: 5);

                // check the return code
                Console.WriteLine($"uninit result = {result}");
                if (result == TIMEOUT_ERROR)
                {
                    throw new Exception("flash uninit timed out");
                }
                else if (result != 0)
                {
                    throw new Exception($"flash uninit failure, code = {result}");
                }
            }

            _activeOperation = FlashOperations.NONE;
        }

        public virtual void PrepareTarget()
        {

        }

        private void EraseAll()
        {
            // 首先检查是否支持整片擦除，如支持则调用整片擦除
            if (!IsEraseAllSupported)
            {
                throw new NotSupportedException("不支持整片擦除");
            }

            // 整片擦除超时时间默认为240秒
            // update core register to execute the erase_all subroutine
            Console.WriteLine("call erase_all");
            var result = CallFunctonAndWait(_flashAlgo.PcEraseAll, timeout: 30);

            // check the return code
            Console.WriteLine($"erase_all result = {result}");
            if (result == TIMEOUT_ERROR)
            {
                throw new Exception("flash uninit timed out");
            }
            else if (result != 0)
            {
                throw new Exception($"flash uninit failure, code = {result}");
            }
        }

        public void EraseSector(FlashRegion flash, uint addr)
        {
            if (addr % flash.SectorSize != 0)
            {
                return;
            }

            var result = CallFunctonAndWait(_flashAlgo.PcEraseSector, addr, timeout: 5);

            // check the return code
            if (result == TIMEOUT_ERROR)
            {
                throw new Exception("flash uninit timed out");
            }
            else if (result != 0)
            {
                throw new Exception($"flash uninit failure, code = {result}");
            }
        }

        private int ProgramPage(byte[] sectorBuffer, uint addr)
        {
            // 加载到内存
            var writeBuf = sectorBuffer;
            _cpu.Ap.WriteBlockMemoryUnaligned8(_flashAlgo.BeginData, writeBuf);

            // 编程到flash
            int result = CallFunctonAndWait(_flashAlgo.PcProgramPage, addr, (uint)writeBuf.Length, _flashAlgo.BeginData, timeout: 5);
            if (result == TIMEOUT_ERROR)
            {
                throw new Exception("flash program timed out");
            }
            else if (result != 0)
            {
                throw new Exception($"flash program failure, code = {result}");
            }
            return result;
        }

        private void EraseSectors(FlashRegion flash, uint startAddr, int length)
        {
            // 计算sector总数
            int sectors = (int)((length + flash.SectorSize - 1) / flash.SectorSize);

            uint curAddr = startAddr;

            for (int i = 0; i < sectors; i++)
            {
                EraseSector(flash, curAddr);
                curAddr += flash.SectorSize;
            }
        }

        public void EraseChip()
        {
            var flash = _target.MemoryMap.Regions.Where(x => x.MemoryType == MemoryType.FLASH).FirstOrDefault() as FlashRegion;

            if (flash == null)
            {
                throw new ArgumentException("未找到flash区域");
            }

            if (IsEraseAllSupported)
            {
                EraseAll();
            }
            else
            {
                EraseSectors(flash, flash.Start, (int)flash.Length);
            }
        }

        public void Program(byte[] buffer, uint addr, FlashEraseModes eraseMode = FlashEraseModes.SECTOR)
        {
            int len = buffer.Length;

            var flashFirst = _target.MemoryMap.Regions.Where(x => x.MemoryType == MemoryType.FLASH).FirstOrDefault();

            if (flashFirst == null)
            {
                throw new ArgumentException("未找到flash区域");
            }

            var flash = flashFirst as FlashRegion;

            // 擦除flash
            if (eraseMode == FlashEraseModes.CHIP)
            {
                if (IsEraseAllSupported)
                {
                    EraseAll();
                }
                else
                {
                    EraseSectors(flash, addr, buffer.Length);
                }
            }
            else if (eraseMode == FlashEraseModes.SECTOR)
            {
                EraseSectors(flash, addr, buffer.Length);
            }
            else
            {
                throw new NotSupportedException();
            }

            int pageSize = (int)_flashAlgo.PageSize;
            // 计算page
            int pages = (len + pageSize - 1) / pageSize;

            uint curAddr = addr;
            int totalBytes = buffer.Length;
            int programmedBytes = 0;

            for (int i = 0; i < pages; i++)
            {
                var sectorBuffer = buffer.Skip(i * pageSize).Take(pageSize).ToArray();
                if (sectorBuffer.Length % _flashAlgo.MinProgramLength != 0)
                {
                    var alignedBuffer = new byte[(sectorBuffer.Length + _flashAlgo.MinProgramLength - 1) / _flashAlgo.MinProgramLength * _flashAlgo.MinProgramLength];
                    Array.Copy(sectorBuffer, alignedBuffer, sectorBuffer.Length);
                    sectorBuffer = alignedBuffer;
                }

                ProgramPage(sectorBuffer, curAddr);
                // Console.WriteLine($"program addr=0x{curAddr:X8} success");
                curAddr += (uint)pageSize;

                // 更新进度
                programmedBytes += pageSize;
                ProgressCallback?.Invoke(programmedBytes, totalBytes);
            }

            // 编程完成后进行校验
            AppendLog("开始校验数据...");
            UpdateDownloadProgress(0);  // 重置进度条

            try
            {
                int verifyPageSize = (int)_flashAlgo.PageSize;
                int verifyPages = (buffer.Length + verifyPageSize - 1) / verifyPageSize;
                uint verifyAddr = addr;
                int verifiedBytes = 0;

                for (int i = 0; i < verifyPages; i++)
                {
                    // 计算当前页的大小
                    int currentPageSize = Math.Min(verifyPageSize, buffer.Length - i * verifyPageSize);
                    int words = (currentPageSize + 3) / 4;  // 向上取整到最接近的4字节

                    // 读取当前页
                    var readbackWords = _cpu.ReadMemoryBlock32(verifyAddr, words);
                    var readbackData = CommonUtils.U32leListToByteList(readbackWords);

                    // 获取原始数据的当前页
                    var originalData = buffer.Skip(i * verifyPageSize).Take(currentPageSize).ToArray();

                    // 比较数据
                    for (int j = 0; j < currentPageSize; j++)
                    {
                        if (originalData[j] != readbackData[j])
                        {
                            uint mismatchAddr = verifyAddr + (uint)j;
                            AppendLog($"数据不匹配 地址: 0x{mismatchAddr:X8}, 期望值: 0x{originalData[j]:X2}, 实际值: 0x{readbackData[j]:X2}");
                            throw new Exception($"数据校验失败，地址 0x{mismatchAddr:X8} 处数据不匹配");
                        }
                    }

                    verifyAddr += (uint)currentPageSize;
                    verifiedBytes += currentPageSize;

                    // 更新进度
                    double progress = (double)verifiedBytes / buffer.Length * 100;
                    UpdateDownloadProgress(progress);
                }

                AppendLog("数据校验成功");
                UpdateDownloadProgress(100);  // 确保进度条到达100%
            }
            catch (Exception ex)
            {
                AppendLog($"校验过程出错: {ex.Message}");
                UpdateDownloadProgress(0);  // 出错时重置进度条
                throw;
            }
        }

        private void UpdateDownloadProgress(double progress)
        {
            ProgressCallback?.Invoke((int)progress, 100);
        }

        public void Cleanup()
        {
            UnInit();
            RestoreTarget();
            _didPrepareTarget = false;
        }

        public virtual void RestoreTarget()
        {
            // pass
        }

        private int CallFunctonAndWait(uint pc, uint? r0 = null, uint? r1 = null, uint? r2 = null, uint? r3 = null, bool init = false, int? timeout = null)
        {
            CallFunction(pc, r0, r1, r2, r3, init);
            return WaitForCompletion(timeout);
        }

        /// <summary>
        /// Wait until the breakpoint is hit.
        /// Checks for:
        /// - Timeout, using the _timeout_ parameter.
        /// - The target is halted after executing the flash operation.
        /// - Stack overflow.
        /// </summary>
        /// <param name="timeout"></param>
        /// <returns></returns>
        private int WaitForCompletion(int? timeout)
        {
            var state = CortexM.TARGET_RUNNING;

            var timeoutInSeconds = new TimeSpan(0, 0, timeout ?? 5);
            var startTime = DateTime.Now;

            DateTime now = DateTime.Now;
            bool isTimeout = true;

            do
            {
                now = DateTime.Now;

                try
                {
                    state = _cpu.GetState();
                    if (state != CortexM.TARGET_RUNNING)
                    {
                        isTimeout = false;
                        break;
                    }
                }
                catch (DAPAccessIntf.TransferTimeoutError)
                {
                    Console.WriteLine("target.get_state probe timeout");
                }
                catch (DAPAccessIntf.TransferFaultError)
                {
                    Console.WriteLine("target.get_state probe fault");
                }
            }
            while (now - startTime < timeoutInSeconds);

            if (isTimeout)
            {
                _cpu.Halt();
                var ipsr = _cpu.ReadCoreRegister("ipsr");
                Console.WriteLine($"flash operation timed out; IPSR=0x{ipsr.IntValue:x8}");
                return TIMEOUT_ERROR;
            }

            // TODO flash_algo_debug
            if (state != CortexM.TARGET_HALTED)
            {
                _cpu.Halt();
                var ipsr = _cpu.ReadCoreRegister("ipsr");
                throw new Exception($"target was not halted as expected after calling flash algorithm routine (IPSR=0x{ipsr.IntValue:x8})");
            }

            // TODO Check stack canary if we have one.

            return (int)_cpu.ReadCoreRegister("r0").IntValue;
        }

        private void CallFunction(uint pc, uint? r0 = null, uint? r1 = null, uint? r2 = null, uint? r3 = null, bool init = false)
        {
            List<string> regList = [];
            List<uint> dataList = [];

            // TODO flash_algo_debug
            regList.Add("pc");
            dataList.Add(pc);

            if (r0 != null)
            {
                regList.Add("r0");
                dataList.Add(r0.Value);
            }

            if (r1 != null)
            {
                regList.Add("r1");
                dataList.Add(r1.Value);
            }

            if (r2 != null)
            {
                regList.Add("r2");
                dataList.Add(r2.Value);
            }

            if (r3 != null)
            {
                regList.Add("r3");
                dataList.Add(r3.Value);
            }

            if (init)
            {
                regList.Add("r9");
                dataList.Add(_flashAlgo.StaticBase);
            }

            if (init)
            {
                regList.Add("sp");
                dataList.Add(_flashAlgo.BeginStack);
            }

            regList.Add("lr");
            dataList.Add(_flashAlgo.LoadAddress + 1);

            _cpu.WriteCoreRegistersRaw([.. regList], [.. dataList]);

            // resume target
            _cpu.Resume();
        }
    }
}
